from typing import Dict, List, Tuple

from annotated_types import Gt, Lt
from pydantic import BaseModel, BeforeValidator
from typing_extensions import Annotated, Literal


def list_constraints(values: list = []):

    def wrapper(v):
        for i in v:
            if i not in values:
                raise ValueError(f'invalid value: {i}')
        return v

    return wrapper


def dict_constraints(keys: list = []):

    def wrapper(v):
        for row in v:
            for i in row:
                if i not in keys:
                    raise ValueError(f'invalid value: {i}')
        return v

    return wrapper


class Fruit(BaseModel):
    name: str
    color: Literal['red', 'green']
    weight: Annotated[float, Gt(2)]
    bazam: Dict[str, List[Tuple[int, bool, float]]]

    # inputNumber 限制取值范围
    age: Annotated[int, Gt(0), Lt(100)]

    # 特定范围内单选
    task_mode: Literal['cls', 'reg']

    # 特定范围内多选，
    metrics: Annotated[List[str],
                       BeforeValidator(list_constraints(['ks', 'acc'])),
                       '评估方法']

    # 特定keys, eg:分层采样>添加采样策略
    custom_ratios: Annotated[
        List[dict],
        BeforeValidator(dict_constraints(['labelVal', 'ratio']))]

    # 文件上传
    upload_model: Annotated[str, 'upload', '模型文件上传']

    # switch开关
    is_exist: bool

    # 生成随机数
    random_number: Annotated[int, Gt(0), 'random']

    # binning_strategies
    binning_strategies: Annotated[Dict[str, int], 'key=column_name,value=1']

    metrics_cls: List[Literal['acc', 'auc', 'f1score']]


print(
    Fruit(
        name='Apple',
        color='red',
        weight=4.2,
        bazam={'foobar': [(1, True, 0.1)]},
        age=10,
        task_mode='cls',
        metrics=['ks'],
        custom_ratios=[{
            'labelVal': 1,
            'ratio': 0.7
        }, {
            'labelVal': 1,
            'ratio': 0.7
        }],
        upload_model='/root/main.py',
        is_exist=True,
        random_number=123,
        binning_strategies={
            'x1': 2,
            'x3': 3
        },
        metrics_cls=['acc', 'f1score']))
